import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader


class MovieRating(Dataset):
    def __init__(self, data:pd.DataFrame):
        self.data = data
        self.max_genre_size = 8

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        df = self.data
        item = df.iloc[idx]
        return item_process(item, self.max_genre_size)



def item_process(item:pd.DataFrame, max_genre_size=8):
    user_id = item['user_id']
    gender = item['gender']
    age = item['age']
    occupation = item['occupation']
    zip_code = item['ZIP']

    user_info = {
        'user_id': torch.tensor(user_id, dtype=torch.long),
        'gender': torch.tensor(gender, dtype=torch.long),
        'age': torch.tensor(age, dtype=torch.long),
        'occupation': torch.tensor(occupation, dtype=torch.long),
        'ZIP': torch.tensor(zip_code, dtype=torch.long)
    }

    movie_id = item['movie_id']
    name = item['name']
    genre = item['genre']
    if len(genre) < max_genre_size:
        genre += [0] * (max_genre_size - len(genre))
    else:
        genre = genre[:max_genre_size]

    movie_info = {
        'movie_id': torch.tensor(movie_id, dtype=torch.long),
        'name': name,
        'genre': torch.tensor(genre, dtype=torch.long)
    }
    rating = torch.tensor(item['rating'], dtype=torch.float)

    return {
        'user_info': user_info,
        'movie_info': movie_info,
        'rating': rating
    }
